/*	
*	File: tcp-header.h
*	Date: 2011/11/28
*	User: Ronmaru
*	Desctiption: tcp header class
*/

#ifndef _TCP_HEADER_H_
#define _TCP_HEADER_H_

//   0                   1                   2                   3   
//   0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 
//   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
//   |          Source Port          |       Destination Port        |
//   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
//   |                        Sequence Number                        |
//   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
//   |                    Acknowledgment Number                      |
//   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
//   |  Data |           |U|A|P|R|S|F|                               |
//   | Offset| Reserved  |R|C|S|S|Y|I|            Window             |
//   |       |           |G|K|H|T|N|N|                               |
//   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
//   |           Checksum            |         Urgent Pointer        |
//   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
//   |                    Options                    |    Padding    |
//   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
//   |                             data                              |
//   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+

const int TCP_MAX_HEADER_LENGTH = 20;

class TCPpseudoheader
{
public:
  TCPpseudoheader() { std::fill(rep_, rep_ + sizeof(rep_), 0); }

  std::string GetHeader() { return std::string(rep_, rep_ + sizeof(rep_)); }

  void SetSourceAddress(boost::asio::ip::address_v4 address) 
  {
    boost::asio::ip::address_v4::bytes_type bytes = address.to_bytes();
    rep_[0] = bytes[0];
    rep_[1] = bytes[1];
    rep_[2] = bytes[2];
    rep_[3] = bytes[3];
  }
  void SetDestinationAddress(boost::asio::ip::address_v4 address)
  {
    boost::asio::ip::address_v4::bytes_type bytes = address.to_bytes();
    rep_[4] = bytes[0];
    rep_[5] = bytes[1];
    rep_[6] = bytes[2];
    rep_[7] = bytes[3];
  }
  void SetMessageProtocol() { Encode(8, 9, 6); }
  void SetDatagramLength(unsigned short n) { Encode(10, 11, n); }

  unsigned int ComputeChecksum() const
  {
    unsigned int sum = 0;

    for (int i = 0; i < 12; i += 2) {
      sum += rep_[i] << 8;
      sum += rep_[i+1];
    }
    return sum;
  }

  boost::asio::ip::address_v4 GetSourceAddress() const
  {
    boost::asio::ip::address_v4::bytes_type bytes
      = { { rep_[0], rep_[1], rep_[2], rep_[3] } };
    return boost::asio::ip::address_v4(bytes);
  }
  boost::asio::ip::address_v4 GetDestinationAddress() const
  {
    boost::asio::ip::address_v4::bytes_type bytes
      = { { rep_[4], rep_[5], rep_[6], rep_[7] } };
    return boost::asio::ip::address_v4(bytes);
  }
  unsigned int GetMessageProtocol() { Decode(8, 9); }
  unsigned int GetDatagramLength() { Decode(10, 11); }

  friend std::istream& operator>>(std::istream& is, TCPpseudoheader& header)
  { return is.read(reinterpret_cast<char*>(header.rep_), 12); }

  friend std::ostream& operator<<(std::ostream& os, const TCPpseudoheader& header)
  { return os.write(reinterpret_cast<const char*>(header.rep_), 12); }

private:
  unsigned short Decode(int a, int b) const
  { return (rep_[a] << 8) + rep_[b]; }

  void Encode(int a, int b, unsigned short n)
  {
    rep_[a] = static_cast<unsigned char>(n >> 8);
    rep_[b] = static_cast<unsigned char>(n & 0xFF);
  }

  unsigned char rep_[12];
};

class TCPheader
{
public:
  TCPheader() { std::fill(rep_, rep_ + sizeof(rep_), 0); }

  void SetSourcePort(unsigned short n) { Encode(0, 1, n); }
  void SetDestinationPort(unsigned short n) { Encode(2, 3, n); }
  void SetSequenceNumber(unsigned long n) { Encode(4, 5, 6, 7, n); }
  void SetAcknoledgmentNumber(unsigned long n) { Encode(8, 9, 10, 11, n); }
  void SetDataOffset(unsigned short n) { rep_[12] = (n << 4); }
  void SetURG(bool n) { rep_[13] = ((~(1 << 5))&rep_[13]) + (n << 5); }
  void SetACK(bool n) { rep_[13] = ((~(1 << 4))&rep_[13]) + (n << 4); }
  void SetPSH(bool n) { rep_[13] = ((~(1 << 3))&rep_[13]) + (n << 3); }
  void SetRST(bool n) { rep_[13] = ((~(1 << 2))&rep_[13]) + (n << 2); }
  void SetSYN(bool n) { rep_[13] = ((~(1 << 1))&rep_[13]) + (n << 1); }
  void SetFIN(bool n) { rep_[13] = ((~(1 << 0))&rep_[13]) + (n << 0); }
  void SetWindow(unsigned short n) { Encode(14, 15, n); }
  void SetChecksum(unsigned short n) { Encode(16, 17, n); }
  void SetUrgentPoint(unsigned short n) { Encode(18, 19, n); }

  unsigned int GetSourcePort() { return Decode(0, 1); }
  unsigned int GetDestinationPort() { return Decode(2, 3); }
  unsigned long GetSequenceNumber() { return Decode(4, 5, 6 ,7); }
  unsigned long GetAcknoledgmentNumber() { return Decode(8, 9, 10, 11); }
  unsigned int GetDataOffset() { return rep_[12] >> 4; }
  bool GetURG() { return !!(rep_[13] & (1 << 5)); }
  bool GetACK() { return !!(rep_[13] & (1 << 4)); }
  bool GetPSH() { return !!(rep_[13] & (1 << 3)); }
  bool GetRST() { return !!(rep_[13] & (1 << 2)); }
  bool GetSYN() { return !!(rep_[13] & (1 << 1)); }
  bool GetFIN() { return !!(rep_[13] & (1 << 0)); }
  unsigned short GetWindow() { return Decode(14, 15); }
  unsigned short GetChecksum() { return Decode(16, 17); }
  unsigned short GetUrgentPoint() { return Decode(18, 19); }

  void ComputeChecksum(const TCPpseudoheader &pseudo_header, std::string &body)
  {
    SetChecksum(0);
    unsigned int sum = 0;

    //pseudo header checksum
    sum += pseudo_header.ComputeChecksum();

    //tcp header checksum
    for (int i = 0; i < TCP_MAX_HEADER_LENGTH; i += 2) {
      sum += rep_[i] << 8;
      sum += rep_[i+1];
    }

    //body checksum
    std::string::iterator body_iter = body.begin();
    while (body_iter != body.end())
    {
      sum += (static_cast<unsigned char>(*body_iter++) << 8);
      if (body_iter != body.end())
        sum += static_cast<unsigned char>(*body_iter++);
    }

    sum = (sum >> 16) + (sum & 0xFFFF);
    sum += (sum >> 16);
    SetChecksum(~sum);
  }

  friend std::istream& operator>>(std::istream& is, TCPheader& header)
  { return is.read(reinterpret_cast<char*>(header.rep_), TCP_MAX_HEADER_LENGTH); }

  friend std::ostream& operator<<(std::ostream& os, const TCPheader& header)
  { return os.write(reinterpret_cast<const char*>(header.rep_), TCP_MAX_HEADER_LENGTH); }

private:
  unsigned short Decode(int a, int b) const
  { return (rep_[a] << 8) + rep_[b]; }
  unsigned long Decode(int a, int b, int c, int d) const
  { return (rep_[a] << 24) + (rep_[b] << 16) + (rep_[c] << 8) + rep_[d]; }

  void Encode(int a, int b, unsigned short n)
  {
    rep_[a] = static_cast<unsigned char>(n >> 8);
    rep_[b] = static_cast<unsigned char>(n & 0xFF);
  }
  void Encode(int a, int b, int c, int d, unsigned long n)
  {
    rep_[a] = static_cast<unsigned char>(n >> 24);
    rep_[b] = static_cast<unsigned char>((n >> 16) & 0xFF);
    rep_[c] = static_cast<unsigned char>((n >> 8) & 0xFF);
    rep_[d] = static_cast<unsigned char>(n & 0xFF);
  }

  unsigned char rep_[TCP_MAX_HEADER_LENGTH];
};


// template <typename Iterator>
// void ComputeChecksum(TCPheader& header, TCPpseudoheader &pseudo_header,
//   Iterator body_begin, Iterator body_end)
// {
//   header.SetChecksum(0);
//   unsigned long int sum = header.GetSourcePort() + header.GetDestinationPort() + header.GetSourcePort();
// 
//   std::string pseudo_header_str = pseudo_header.GetHeader();
//   for (int i = 0; i < pseudo_header_str.length(); i += 2) {
//     sum += static_cast<unsigned int>(pseudo_header_str[i]) << 8;
//     sum += static_cast<unsigned int>(pseudo_header_str[i + 1]);
//   }
// 
//   Iterator body_iter = body_begin;
//   while (body_iter != body_end)
//   {
//     sum += (static_cast<unsigned char>(*body_iter++) << 8);
//     if (body_iter != body_end)
//       sum += static_cast<unsigned char>(*body_iter++);
//   }
// 
//   sum = (sum >> 16) + (sum & 0xFFFF);
//   sum += (sum >> 16);
//   header.SetChecksum(static_cast<unsigned short>(~sum));
// }


#endif //_TCP_HEADER_H_